from dataGeter import inputfomating
import pandas as pd


def getRangeK(df, by="y", col="date"):
    """
    df = pd.read_csv('tables/TradeHistory/000001.SZ.csv')
    df = getRangeK(df, by='Y')
    + by=分组方式
        + Y = year
        + W = week
        + Q = 季度
        + M = 月)
    + col 列的名称 默认 'date'
        + 暂时只适用于 YMD格式不支持时间
        + 如果输入含有 H：M：S格式的数据，请先拆分
    + return: 含['range'] 标签的 Dataframe
    """
    df[["y", "m", "d"]] = df[col].str.split("-", expand=True)
    if by == "Y":  # 年数据
        df["range"] = df["y"]
    elif by == "Q":
        volumns = "Q"
        df["m"] = df["m"].apply(lambda x: int(x))
        df["range"] = df["m"].apply(
            lambda x: "-12-31"
            if x < 4
            else ("-06-30")
            if x < 7
            else ("-09-30" if x < 10 else "-03-31")
        )
        df["range"] = df["y"] + df["range"]
    elif by == "W":
        df["range"] = df["date"].apply(lambda x: "w" + str(TT.getCal(x)))
        df["range"] = df["y"] + df["range"]
    elif by == "M":
        # 获取月数据
        df["range"] = df["y"] + df["m"]
    else:
        print("超出选项")
    return df


def ToRangeK(df):
    """
    df = pd.read_csv('tables/TradeHistory/000001.SZ.csv')
    df = getRangeK(df, by='Y')
    df = ToRangeK(df)
    通过getRangeK调用获取含有'range'的列表
    获取分组后的子列表
    """

    def dataChack(n, name, types):
        if n != 0:
            return n
        else:
            print(n, name, types)
            return n

    nl = df.columns.tolist()
    df = df.groupby("range")
    lists = []
    for name, group in df:
        dic = {}
        dic["date"] = name
        dic["high"] = dataChack(group["high"].max(), group, "high")
        dic["low"] = dataChack(group["low"].min(), group, "low")
        dic["open"] = dataChack(group.iloc[0]["open"], group, "open")
        dic["close"] = dataChack(group.iloc[-1]["close"], group, "close")
        if "volume" in nl:
            dic["volume"] = dataChack(group["volume"].sum().round(5), group, "volume")
        elif "Volume" in nl:
            dic["Volume"] = dataChack(group["Volume"].sum().round(5), group, "volume")
        lists.append(dic)
    df = pd.DataFrame(lists)
    df = df.set_index("date")
    return df


def SortBYINC(df, code, by="Q", col='date'):
    """
    导入df,含股K线数据票信息，
    调用getRangeK，ToRangeK处理数据。
    输出含区间内涨跌幅的值
    code = code
    """
    df = getRangeK(df, by=by, col=col)
    df = ToRangeK(df)
    df["code"] = code
    df["stictINC"] = df["close"] - df["open"]
    df["maxINC"] = df["high"] - df["low"]
    df["stictINC%"] = df["stictINC"] / df["open"]
    df["maxINC%"] = df["maxINC"] / df["low"]
    return df